#!/usr/bin/python
# This program prints Pfam domains in order. 
# Works pretty good right now but method to remove overlaps is still rough.. Need better algorithm to keep more relevant one.
# Example, maybe use E value or overlap length as criteria to keep the best hit among the overlaps.
# Last edit: 04-13-2011

import sys
import re
from Bio.Blast import NCBIXML

xml_file = sys.argv[1]
xf = open(xml_file, "rU")
r = NCBIXML.parse(xf)
results_rec = r.next()

m1 = re.compile('.* (COG\d+), (\w+),.*')
m2 = re.compile('(COG\d+), (\w+|\w+-\w+|\w+-\w+-\w+|\{.*\}),.*') #2nd group matches regular words or hyphenated words
m3 = re.compile('(.*).cog.rpsblast.xml')

desc_obj_list = []


cds_id = m3.match(xml_file).group(1)


max_evalue = 0.01 #0.01 matches well with SMART EMBL's domain search results

pfam_ids = []
pfam_descs = []
query_starts = []
query_stops = []
evalues = []

if len(results_rec.alignments) == 0:
    #print cds_id + "\t" + "No Pfam Hit"
    pass
else:
#    for rec in results_rec.descriptions:
#        if rec.e < 0.01:
#            desc_obj_list.append(rec)
#            p = m1.match(rec.title)
#            for a in results_rec.alignments:
#                for hsp in a.hsps:
#                    print rec.num_alignments, rec.score, rec.e, p.group(1), p.group(2)
#    print "Protein length:" + "\t" + str(results_rec.query_letters)
#    print "# of alignments:" + "\t" + str(len(results_rec.alignments))

    for alignment in results_rec.alignments:
        p = m2.match(alignment.hit_def)
        pfam_id = p.group(1)
        pfam_desc = p.group(2)
        if len(alignment.hsps) == 1:
            for hsp in alignment.hsps:
                if hsp.expect < max_evalue:
                    #print cds_id + "\t" + pfam_id + "\t" + pfam_desc + "\t" + str(hsp.query_start) + "\t" + str(hsp.query_end) + "\t" +str(hsp.expect)
                    pfam_ids.append(pfam_id)
                    pfam_descs.append(pfam_desc)
                    query_starts.append(hsp.query_start)
                    query_stops.append(hsp.query_end)
                    evalues.append(hsp.expect)
        else:
            hsps_list = []
            for hsp in alignment.hsps:
                hsps_list.append(hsp)
#            if hsps_list[0].expect < 0.01:
#                pfam_ids.append(pfam_id)
#                pfam_descs.append(pfam_desc)
#                query_starts.append(hsps_list[0].query_start)
#                query_stops.append(hsps_list[0].query_end)
#                evalues.append(hsps_list[0].expect)
            for h in hsps_list:
                if h.expect < max_evalue:
                    pfam_ids.append(pfam_id)
                    pfam_descs.append(pfam_desc)
                    query_starts.append(h.query_start)
                    query_stops.append(h.query_end)
                    evalues.append(h.expect)

i = 0
num = len(pfam_ids)
final_pids = []
final_domains = []
final_starts = []
final_stops = []
final_coords = []
final_evalues = []

while i < num:
    if i == 0:
        final_pids.append(pfam_ids[i])
        final_domains.append(pfam_descs[i])
        final_starts.append(query_starts[i])
        final_stops.append(query_stops[i])
        final_coords.append((query_starts[i], query_stops[i]))
        final_evalues.append(evalues[i])
    if i == num - 1:
        pass
    else:
        if query_starts[i] < query_starts[i+1] < query_stops[i] or query_starts[i] < query_stops[i+1] < query_stops[i]:
            pass
        elif query_starts[i] == query_starts[i+1] or query_stops[i] == query_stops[i+1]:
            pass
        else:
            final_pids.append(pfam_ids[i+1])
            final_domains.append(pfam_descs[i+1])
            final_starts.append(query_starts[i+1])
            final_stops.append(query_stops[i+1])
            final_coords.append((query_starts[i+1], query_stops[i+1]))
            final_evalues.append(evalues[i+1])
    i += 1

#for i, j in enumerate(final_pids):
#    print "Final", cds_id + "\t" + final_pids[i] + "\t" + final_domains[i] + "\t" + str(final_starts[i]) \
#        + "\t" + str(final_stops[i]) + "\t" + str(final_evalues[i])

sorted_pids = []
sorted_stops = []
sorted_evalues = []
#--------
#Sorting the starts. Can remove this part since code is effectively replaced by tuple-based coordinates code below.
sorted_starts = []
sorted_domains = []
domains_dict = {}
for i, j in enumerate(final_starts):
    domains_dict[j] = final_domains[i] #to create dictionary of domains with keys: '0' -> 'start'

sorted_starts = sorted(final_starts) #sort the starts to order starting domain

for x in sorted_starts:
    sorted_domains.append(domains_dict[x]) #add the domain one by one in order (hopefully..)
#--------
#Sorting the coordinates. Keep this part.
#fd = []
sd = []
dd = {}
ed = {}

for i, j in enumerate(final_coords):
    dd[j] = final_domains[i] #to create dictionary of domains with keys to their coordinate tuples: (start, stop) -> 'AAA'
    ed[j] = final_evalues[i] #dict of evalues with coordinates: (start, stop) -> '1.3e-7'

sorted_coords = sorted(final_coords) #sort the coordinate tuples to order first to last domains
#--------

#Check if any coords are overlapping. Now they should be sorted next to each other for checking
#Eg: ((1, 50), (2, 30), (56, 90), (60, 100), etc...)
z = 0
num = len(sorted_coords)

#for i, j in enumerate(dd):
#    print dd[j], j

while z < num:
    if z == 0:
        sd.append(dd[sorted_coords[z]])
    if z == num - 1:
        pass
        #sd.append(dd[sorted_coords[z]])
    else:
        # x1 < x2 < y1 or x1 < y2 < y1
        if sorted_coords[z][0] < sorted_coords[z+1][0] < sorted_coords[z][1] or sorted_coords[z][0] < sorted_coords[z+1][1] < sorted_coords[z][1]:
            pass
            #if ed[sorted_coords[z]] < ed[sorted_coords[z+1]]:
            #    sd.append(dd[sorted_coords[z]])
            #else:
            #    sd.append(dd[sorted_coords[z+1]])
        # x1 == x2 or y1 == y2
        elif sorted_coords[z][0] == sorted_coords[z+1][0] or sorted_coords[z][1] == sorted_coords[z+1][1]:
            pass
            #if ed[sorted_coords[z]] < ed[sorted_coords[z+1]]:
            #    sd.append(dd[sorted_coords[z]])
            #else:
            #    sd.append(dd[sorted_coords[z+1]])
        # x1 < x2 < y2 < y1
        elif sorted_coords[z][0] < sorted_coords[z+1][0] < sorted_coords[z][1] and sorted_coords[z][0] < sorted_coords[z+1][1] < sorted_coords[z][1]:
            pass
            #if ed[sorted_coords[z]] < ed[sorted_coords[z+1]]:
            #    sd.append(dd[sorted_coords[z]])
            #else:
            #    sd.append(dd[sorted_coords[z+1]])
        # x1 == x2 and y1 == y2
        elif sorted_coords[z][0] == sorted_coords[z+1][0] and sorted_coords[z][1] == sorted_coords[z+1][1]:
            pass
            #if ed[sorted_coords[z]] < ed[sorted_coords[z+1]]:
            #    sd.append(dd[sorted_coords[z]])
            #else:
            #    sd.append(dd[sorted_coords[z+1]])
        else:
            sd.append(dd[sorted_coords[z+1]])
    z += 1

#print "Domains:"
if len(results_rec.alignments) == 0:
    #print cds_id + "\t" + "--No Pfam Hit--"
    print cds_id + "\t" + "----"
else:
    #print cds_id + "\t" + "; ".join(final_domains)
    #print cds_id + "\t" + ", ".join(sorted_domains)
    print cds_id + "\t" + " + ".join(sd)

xf.close()